//! Span管理模块
//! 
//! 提供Span的创建、管理和操作工具

use crate::*;
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

impl TraceSpan {
    /// 创建新的根Span
    pub fn new_root(operation_name: &str, trace_id: String, span_id: String) -> Self {
        Self {
            trace_id,
            span_id,
            parent_span_id: None,
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            finish_time: None,
            tags: HashMap::new(),
            logs: Vec::new(),
        }
    }

    /// 创建子Span
    pub fn new_child(&self, operation_name: &str, span_id: String) -> Self {
        Self {
            trace_id: self.trace_id.clone(),
            span_id,
            parent_span_id: Some(self.span_id.clone()),
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            finish_time: None,
            tags: HashMap::new(),
            logs: Vec::new(),
        }
    }

    /// 检查Span是否已完成
    pub fn is_finished(&self) -> bool {
        self.finish_time.is_some()
    }

    /// 获取Span持续时间（毫秒）
    pub fn duration_ms(&self) -> Option<f64> {
        self.finish_time.map(|finish_time| {
            if let (Ok(start), Ok(finish)) = (
                self.start_time.duration_since(UNIX_EPOCH),
                finish_time.duration_since(UNIX_EPOCH)
            ) {
                (finish.as_millis() - start.as_millis()) as f64
            } else {
                0.0
            }
        })
    }

    /// 添加标签
    pub fn add_tag(&mut self, key: &str, value: &str) {
        self.tags.insert(key.to_string(), value.to_string());
    }

    /// 添加多个标签
    pub fn add_tags(&mut self, tags: HashMap<String, String>) {
        self.tags.extend(tags);
    }

    /// 获取标签值
    pub fn get_tag(&self, key: &str) -> Option<&String> {
        self.tags.get(key)
    }

    /// 添加日志
    pub fn add_log(&mut self, message: &str) {
        self.add_log_with_fields(message, HashMap::new());
    }

    /// 添加带字段的日志
    pub fn add_log_with_fields(&mut self, message: &str, fields: HashMap<String, String>) {
        let log = SpanLog {
            timestamp: SystemTime::now(),
            message: message.to_string(),
            fields,
        };
        self.logs.push(log);
    }

    /// 设置错误标签
    pub fn set_error(&mut self, error: &str) {
        self.add_tag("error", "true");
        self.add_tag("error.message", error);
    }

    /// 设置HTTP相关标签
    pub fn set_http_info(&mut self, method: &str, url: &str, status_code: u16) {
        self.add_tag("http.method", method);
        self.add_tag("http.url", url);
        self.add_tag("http.status_code", &status_code.to_string());
    }

    /// 设置数据库相关标签
    pub fn set_db_info(&mut self, db_type: &str, db_name: &str, statement: &str) {
        self.add_tag("db.type", db_type);
        self.add_tag("db.name", db_name);
        self.add_tag("db.statement", statement);
    }

    /// 设置服务信息
    pub fn set_service_info(&mut self, service_name: &str, service_version: &str) {
        self.add_tag("service.name", service_name);
        self.add_tag("service.version", service_version);
    }

    /// 完成Span
    pub fn finish(&mut self) {
        if self.finish_time.is_none() {
            self.finish_time = Some(SystemTime::now());
        }
    }

    /// 获取Span上下文信息（用于跨服务传递）
    pub fn get_context(&self) -> SpanContext {
        SpanContext {
            trace_id: self.trace_id.clone(),
            span_id: self.span_id.clone(),
            baggage: self.tags.clone(),
        }
    }
}

impl SpanLog {
    /// 创建新的日志条目
    pub fn new(message: &str) -> Self {
        Self {
            timestamp: SystemTime::now(),
            message: message.to_string(),
            fields: HashMap::new(),
        }
    }

    /// 创建带字段的日志条目
    pub fn with_fields(message: &str, fields: HashMap<String, String>) -> Self {
        Self {
            timestamp: SystemTime::now(),
            message: message.to_string(),
            fields,
        }
    }

    /// 添加字段
    pub fn add_field(&mut self, key: &str, value: &str) {
        self.fields.insert(key.to_string(), value.to_string());
    }
}

/// Span上下文，用于跨服务传递
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpanContext {
    pub trace_id: String,
    pub span_id: String,
    pub baggage: HashMap<String, String>,
}

impl SpanContext {
    /// 创建新的Span上下文
    pub fn new(trace_id: String, span_id: String) -> Self {
        Self {
            trace_id,
            span_id,
            baggage: HashMap::new(),
        }
    }

    /// 添加行李（baggage）数据
    pub fn add_baggage(&mut self, key: &str, value: &str) {
        self.baggage.insert(key.to_string(), value.to_string());
    }

    /// 获取行李数据
    pub fn get_baggage(&self, key: &str) -> Option<&String> {
        self.baggage.get(key)
    }

    /// 序列化为HTTP头格式
    pub fn to_headers(&self) -> HashMap<String, String> {
        let mut headers = HashMap::new();
        headers.insert("x-trace-id".to_string(), self.trace_id.clone());
        headers.insert("x-span-id".to_string(), self.span_id.clone());
        
        for (key, value) in &self.baggage {
            headers.insert(format!("x-baggage-{}", key), value.clone());
        }
        
        headers
    }

    /// 从HTTP头反序列化
    pub fn from_headers(headers: &HashMap<String, String>) -> Option<Self> {
        let trace_id = headers.get("x-trace-id")?.clone();
        let span_id = headers.get("x-span-id")?.clone();
        
        let mut baggage = HashMap::new();
        for (key, value) in headers {
            if key.starts_with("x-baggage-") {
                let baggage_key = key.strip_prefix("x-baggage-").unwrap();
                baggage.insert(baggage_key.to_string(), value.clone());
            }
        }
        
        Some(Self {
            trace_id,
            span_id,
            baggage,
        })
    }
}

/// Span构建器，用于方便地创建Span
pub struct SpanBuilder {
    operation_name: String,
    trace_id: Option<String>,
    span_id: Option<String>,
    parent_span_id: Option<String>,
    tags: HashMap<String, String>,
}

impl SpanBuilder {
    /// 创建新的Span构建器
    pub fn new(operation_name: &str) -> Self {
        Self {
            operation_name: operation_name.to_string(),
            trace_id: None,
            span_id: None,
            parent_span_id: None,
            tags: HashMap::new(),
        }
    }

    /// 设置trace_id
    pub fn with_trace_id(mut self, trace_id: String) -> Self {
        self.trace_id = Some(trace_id);
        self
    }

    /// 设置span_id
    pub fn with_span_id(mut self, span_id: String) -> Self {
        self.span_id = Some(span_id);
        self
    }

    /// 设置父Span ID
    pub fn with_parent(mut self, parent_span_id: String) -> Self {
        self.parent_span_id = Some(parent_span_id);
        self
    }

    /// 添加标签
    pub fn with_tag(mut self, key: &str, value: &str) -> Self {
        self.tags.insert(key.to_string(), value.to_string());
        self
    }

    /// 添加多个标签
    pub fn with_tags(mut self, tags: HashMap<String, String>) -> Self {
        self.tags.extend(tags);
        self
    }

    /// 构建Span
    pub fn build(self) -> TracingResult<TraceSpan> {
        let trace_id = self.trace_id.ok_or_else(|| 
            TracingError::InvalidData("trace_id is required".to_string()))?;
        let span_id = self.span_id.ok_or_else(|| 
            TracingError::InvalidData("span_id is required".to_string()))?;

        let span = TraceSpan {
            trace_id,
            span_id,
            parent_span_id: self.parent_span_id,
            operation_name: self.operation_name,
            start_time: SystemTime::now(),
            finish_time: None,
            tags: self.tags,
            logs: Vec::new(),
        };

        Ok(span)
    }
}

/// Span工具函数
pub mod span_utils {
    use super::*;
    
    /// 从多个Span中构建调用链
    pub fn build_trace_tree(spans: &[TraceSpan]) -> Vec<TraceNode> {
        let mut roots = Vec::new();
        let mut span_map: HashMap<String, &TraceSpan> = HashMap::new();
        let mut children_map: HashMap<String, Vec<&TraceSpan>> = HashMap::new();

        // 构建映射
        for span in spans {
            span_map.insert(span.span_id.clone(), span);
            
            if let Some(parent_id) = &span.parent_span_id {
                children_map.entry(parent_id.clone()).or_insert_with(Vec::new).push(span);
            }
        }

        // 找到根节点
        for span in spans {
            if span.parent_span_id.is_none() {
                roots.push(build_node(span, &children_map));
            }
        }

        roots
    }

    fn build_node(span: &TraceSpan, children_map: &HashMap<String, Vec<&TraceSpan>>) -> TraceNode {
        let children = children_map.get(&span.span_id)
            .map(|child_spans| {
                child_spans.iter()
                    .map(|child| build_node(child, children_map))
                    .collect()
            })
            .unwrap_or_default();

        TraceNode {
            span: span.clone(),
            children,
        }
    }

    /// 计算调用链的总持续时间
    pub fn calculate_total_duration(spans: &[TraceSpan]) -> f64 {
        spans.iter()
            .filter_map(|span| span.duration_ms())
            .sum()
    }

    /// 找到最耗时的Span
    pub fn find_slowest_span(spans: &[TraceSpan]) -> Option<&TraceSpan> {
        spans.iter()
            .filter(|span| span.is_finished())
            .max_by(|a, b| {
                a.duration_ms().unwrap_or(0.0)
                    .partial_cmp(&b.duration_ms().unwrap_or(0.0))
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
    }

    /// 按操作名称统计Span
    pub fn group_by_operation(spans: &[TraceSpan]) -> HashMap<String, Vec<&TraceSpan>> {
        let mut groups = HashMap::new();
        
        for span in spans {
            groups.entry(span.operation_name.clone())
                .or_insert_with(Vec::new)
                .push(span);
        }
        
        groups
    }
}

/// 调用链树节点
#[derive(Debug, Clone)]
pub struct TraceNode {
    pub span: TraceSpan,
    pub children: Vec<TraceNode>,
}

impl TraceNode {
    /// 获取节点深度
    pub fn depth(&self) -> usize {
        if self.children.is_empty() {
            1
        } else {
            1 + self.children.iter().map(|child| child.depth()).max().unwrap_or(0)
        }
    }

    /// 获取所有后代Span
    pub fn all_spans(&self) -> Vec<&TraceSpan> {
        let mut spans = vec![&self.span];
        for child in &self.children {
            spans.extend(child.all_spans());
        }
        spans
    }
}